package com.xykj.deploy.config;

import com.jcraft.jsch.ChannelShell;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.xykj.deploy.host.entity.Host;
import com.xykj.deploy.host.service.IHostService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.Executors;

/**
 * @author Ethan
 * @version 1.0
 * @description: TODO
 * @date 2024-12-02 14:33
 */
@Slf4j
@Component
public class WebSSHHandler extends TextWebSocketHandler {

    @Resource
    private IHostService hostService;

    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception {
        String hostIp = getQueryParam(webSocketSession, "hostIp");
        if (hostIp == null) {
            webSocketSession.close(CloseStatus.BAD_DATA);
            return;
        }

        JSch jsch = new JSch();
        try {
            // 替换为实际用户名和密码
            Host host = hostService.getByIp(hostIp);
            Session session = jsch.getSession(host.getUsername(), hostIp, Integer.valueOf(host.getPort()));
            session.setPassword(host.getPassword());
            session.setConfig("StrictHostKeyChecking", "no");

            log.info("Connecting to SSH To Host: {}", hostIp);
            session.connect(60000);
            log.info("Connected to Host: {}", hostIp);

            ChannelShell channel = (ChannelShell) session.openChannel("shell");
            channel.setPty(true);

            InputStream inputStream = channel.getInputStream();
            OutputStream outputStream = channel.getOutputStream();

            channel.connect(3000);

            webSocketSession.getAttributes().put(hostIp, new SSHSession(session, channel, outputStream));

            Executors.newCachedThreadPool().submit(() -> {
                byte[] buffer = new byte[1024];
                try {
                    int length;
                    while ((length = inputStream.read(buffer)) != -1) {
                        if (!webSocketSession.isOpen()) {
                            break;
                        }
                        webSocketSession.sendMessage(new TextMessage(new String(buffer, 0, length, StandardCharsets.UTF_8)));
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                }
            });
        } catch (JSchException e) {
            log.error("SSH Connection Error: {}", e.getMessage());
            webSocketSession.close(CloseStatus.SERVER_ERROR);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession webSocketSession, TextMessage message) throws Exception {
        String hostIp = getQueryParam(webSocketSession, "hostIp");
        SSHSession sshSession = (SSHSession) webSocketSession.getAttributes().get(hostIp);
        if (sshSession != null) {
            sshSession.outputStream().write(message.getPayload().getBytes(StandardCharsets.UTF_8));
            sshSession.outputStream().flush();
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus status) throws Exception {
        String hostIp = getQueryParam(webSocketSession, "hostIp");
        SSHSession sshSession = (SSHSession) webSocketSession.getAttributes().get(hostIp);
        if (sshSession != null) {
            log.info("DisConnected to Host: {}", hostIp);
            sshSession.close();
        }
    }

    private String getQueryParam(WebSocketSession session, String paramName) {
        String query = session.getUri().getQuery();
        if (query != null) {
            for (String param : query.split("&")) {
                String[] pair = param.split("=");
                if (pair.length == 2 && pair[0].equals(paramName)) {
                    return pair[1];
                }
            }
        }
        return null;
    }

}